import os
import os.path as osp
import math
import collections
import cv2 as cv
import numpy as np
from shapely.geometry import Polygon, MultiPoint  # 多边形


IMG_SIZE = 3000

BASEDIR = osp.abspath('.')
path_results_txt = osp.join(BASEDIR, 'results-conf0.1-best-BiCnnNew.txt')
gt_labels_dir = osp.join(BASEDIR, 'val_txt')

correct_dict = collections.defaultdict(int)
over_alert_dict = collections.defaultdict(int)
gt_dict = collections.defaultdict(int)


def compute_one_ration_iou(x_box, y_box):
    # 四边形的二维坐标表示
    xx_box, yy_box = np.array(x_box).reshape(4, 2), np.array(y_box).reshape(4, 2)
    # 构建四边形对象，会自动计算四个点的顺序：左上 左下  右下 右上 左上（返回5个点，最后回到起始点）
    x_poly, y_poly = Polygon(xx_box).convex_hull, Polygon(yy_box).convex_hull

    intersect_area = x_poly.intersection(y_poly).area  # 相交面积
    if intersect_area == 0:
        iou = 0
    else:
        union_area = x_poly.area + y_poly.area - intersect_area  # 总共面积
        iou = intersect_area / union_area
    return iou


def convert_to_xy4(x, y, w, h, angle):
    arg = ((x, y), (w, h), angle)
    box = cv.boxPoints(arg)
    return np.int0(box)


def is_matched(gt_xy4, pred_xy4):
    if compute_one_ration_iou(gt_xy4, pred_xy4) > 0.1:
        return True
    else:
        return False


def calculate_over_alert():
    with open(path_results_txt, 'r') as fp:
        lines = fp.readlines()

    for line in lines:  # for each predicted box
        flag = False  # 初始假定为虚警框
        one_box_info = line.strip('\n').split(' ')
        path_img, cls_name, conf, *pred_xy4 = one_box_info

        path_txt = osp.join(gt_labels_dir, path_img.split('/')[-1].split('.')[0]+'.txt')
        with open(path_txt, 'r') as fp:
            gt_infos = fp.readlines()
        # print(gt_infos)  # 该预测框所在图中所有Ground Truth框

        for gt_info in gt_infos:  # for each ground truth box in this picture
            cls0, x, y, w, h, angle = gt_info.strip('\n').split(' ')
            x = eval(x) * IMG_SIZE
            y = eval(y) * IMG_SIZE
            w = eval(w) * IMG_SIZE
            h = eval(h) * IMG_SIZE
            angle = int(eval(angle) * 180 / math.pi)

            # 都转换成四顶点形式
            gt_xy4 = convert_to_xy4(x, y, w, h, angle)
            # print('gt_xy4:', gt_xy4)

            # Ground Truth是否与检测框匹配
            if is_matched(gt_xy4, pred_xy4):
                flag = True  # 判断为正确框
                break

        # 最后判断一下，如果这个是正确框，
        if flag is True:
            correct_dict[path_img.split('/')[-1].split('.')[0]]+=1

        # 如果这个是虚警框
        else:
            over_alert_dict[path_img.split('/')[-1].split('.')[0]]+=1


calculate_over_alert()
print('\ncorrect:\n', correct_dict)
print('\nover_alert_dict:\n', over_alert_dict)
